package demo.kk;

import java.util.Arrays;
import java.util.Random;

public class TopKElements {

    public static int[] findTopK(int[] arr, int k) {
        if (k <= 0) return new int[0];
        if (k >= arr.length) return arr.clone();

        int[] copy = arr.clone(); // 避免修改原数组
        quickSelect(copy, 0, copy.length - 1, copy.length - k);
        return Arrays.copyOfRange(copy, copy.length - k, copy.length);
    }

    private static void quickSelect(int[] arr, int low, int high, int targetIndex) {
        if (low >= high) return;

        int pivotIndex = randomPartition(arr, low, high);
        if (pivotIndex == targetIndex) {
            return;
        } else if (pivotIndex < targetIndex) {
            quickSelect(arr, pivotIndex + 1, high, targetIndex);
        } else {
            quickSelect(arr, low, pivotIndex - 1, targetIndex);
        }
    }

    private static int randomPartition(int[] arr, int low, int high) {
        // 随机选择轴点以避免最坏情况
        int randomIndex = new Random().nextInt(high - low + 1) + low;
        swap(arr, low, randomIndex);
        
        int key = arr[low];
        int hole = low;
        int left = low;
        int right = high;

        while (left < right) {
            while (left < right && arr[right] >= key) {
                right--;
            }
            arr[hole] = arr[right];
            hole = right;

            while (left < right && arr[left] <= key) {
                left++;
            }
            arr[hole] = arr[left];
            hole = left;
        }
        arr[hole] = key;
        return hole;
    }

    private static void swap(int[] arr, int i, int j) {
        int temp = arr[i];
        arr[i] = arr[j];
        arr[j] = temp;
    }

    public static void main(String[] args) {
        int[] arr = {5, 4, 3, 2, 1};
        int k = 3;
        int[] topK = findTopK(arr, k);
        System.out.println("Top " + k + " elements: " + Arrays.toString(topK));
    }
}